# flask_app.py
import gc
import random

from flask import Flask, request, jsonify
from waitress import serve
import pandas as pd
import numpy as np
import joblib
import torch

from PathFormer import PathFormer
from predict_utils import predict_time_range

app = Flask(__name__)

# 全局版本变量
version = "v2.1.1"

print("loading global variables...")
# 提取文件路径进行管理
data_path = 'ZGGG.09.02.2023.09.02.2025.1.0.0.cn.utf8.00000000.xls'
encoder_dd_path = f'label_encoder_dd_{version}.pkl'
encoder_c_path = f'label_encoder_c_{version}.pkl'
scaler_path = f'data_scaler_{version}.pkl'
model_dict_path = f'pathformer_dict_{version}.pkl'

print("loading external files...")
# 全局加载模型和编码器（启动时加载一次）
encoder_dd = joblib.load(encoder_dd_path)
encoder_c = joblib.load(encoder_c_path)
scaler = joblib.load(scaler_path)

print("loading data file...")
# 在全局加载一次训练数据文件
data = pd.read_excel(data_path, skiprows=6)

print("initializing model...")
# 初始化模型
model = PathFormer(
    input_dim=15,
    patch_size=48,
    embed_dim=256,
    num_heads=8,
    num_layers=6,
    output_dim=15,
    ff_dim=1024
)

print("loading parameters dictionary...")
# 加载模型参数字典
model.load_state_dict(torch.load(model_dict_path))
# 设备配置并加载到设备上
device = torch.device('cpu')
model.to(device)
model.eval()    # 设置模型为评估模式

print("loading model end...")

def add_time_features(df):
    df['hour_sin'] = np.sin(2 * np.pi * df.index.hour / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df.index.hour / 24)
    df['month_sin'] = np.sin(2 * np.pi * df.index.month / 12)
    df['month_cos'] = np.cos(2 * np.pi * df.index.month / 12)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df.index.dayofweek / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df.index.dayofweek / 7)
    return df


def preprocess_data(original_data: pd.DataFrame) -> pd.DataFrame:
    """数据预处理（需与训练时一致）"""
    processed_data = original_data.drop(columns=['ff10', 'WW', "W'W'"], inplace=False)
    processed_data['当地时间 广州(机场)'] = pd.to_datetime(processed_data['当地时间 广州(机场)'], format='%d.%m.%Y %H:%M')
    processed_data.set_index('当地时间 广州(机场)', inplace=True)
    processed_data.sort_index(ascending=True, inplace=True)

    processed_data = add_time_features(processed_data)

    processed_data['VV'] = processed_data['VV'].astype(str)
    processed_data['VV'] = processed_data['VV'].str.extract(r'(\d+\.\d+)').astype(float)
    processed_data['DD'] = encoder_dd.transform(processed_data[['DD']])
    processed_data['c'] = encoder_c.transform(processed_data[['c']])
    processed_data.ffill(inplace=True)

    return processed_data


def set_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(seed)


@app.route('/predict', methods=['POST'])
def predict():

    # 设置随机种子
    set_seed(42)

    # 解析请求参数
    req = request.json
    start_time = req.get('start_time')
    num_steps = min(int(req.get('num_steps', 10)), 336)  # 限制最大336步

    try:
        # 加载默认训练数据（示例需替换为实际数据源）
        processed_data = preprocess_data(data)

        # 执行预测
        pred_results = predict_time_range(
            model, scaler, processed_data, start_time, num_steps, seq_length=672
        )

        # 反编码结果
        # 将 Series 转换为 DataFrame
        pred_results['DD'] = encoder_dd.inverse_transform(pred_results[['DD']].astype(int)).flatten()
        pred_results['c'] = encoder_c.inverse_transform(pred_results[['c']].astype(int)).flatten()

        # 转换为可序列化的字典
        pred_results.index = pred_results.index.strftime('%Y-%m-%d %H:%M:%S')
        pred_results.rename(columns={pred_results.index.name: 'time'}, inplace=True)
        pred_results = pred_results.iloc[:, :9]

        # 将时间列单独提取出来，并作为每个记录中的一个键值对
        data_list = []
        for index, row in pred_results.iterrows():
            record = row.to_dict()
            record['time'] = index  # 添加时间键值对
            data_list.append(record)

        return jsonify({"status": "success", "data": data_list})

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 400

    finally:
        gc.collect()


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

